# -*- encoding: utf-8 -*-
'''
@File    :   ae.py
@Time    :   2022/3/23 13:59
@Author  :   ZhangChaoYang
@Desc    :   自动编码机AutoEncoder。用正常样本训练AE模型，取拟合误差的99%分位点作为阈值，预期异常样本的拟合误差会大于该阈值
'''

import os
import sys

sys.path.insert(0, os.getcwd())
from tensorflow.keras import Model
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.callbacks import EarlyStopping
from models.encoders import FCEncoder, ConvEncoder
from models.decoders import FCDecoder, ConvDecoder
from models import losses
from models import metrics
from util.view import draw_train_history, draw_fit_error_1d, draw_fit_error_2d, draw_fit_error
from util.err_analyze import fit_err_percentage
from util.work_flow import preprocess, gen_outfile, save_model, load_model
from absl import flags
from absl import app

FLAGS = flags.FLAGS

flags.DEFINE_string('kind', "fc", 'AE模型使用哪种网络，可选fc或conv')
flags.DEFINE_string('corpus', "jiangnan", '数据集，可选cwru、jiangnan、paderborn')
flags.DEFINE_string('data_dim', "1d", '数据集，可选1d、2d')
flags.DEFINE_string('data_trans', "original",
                    '数据集，data_dim=1d时data_trans可选original、fft、stat，data_dim=2d时data_trans可选sfft、cwt')


class AE(Model):
    def __init__(self, kind, input_shape=None, hidden_dims=None, kernel_size=[], strides=[], filters=[]):
        super(AE, self).__init__()
        if kind == "fc":
            self.encoder = FCEncoder(hidden_dims)
            self.decoder = FCDecoder(input_shape, hidden_dims)
        elif kind == "conv":
            height, width = input_shape
            self.encoder = ConvEncoder(kernel_size, strides, filters)
            self.decoder = ConvDecoder(height, width, kernel_size, strides, filters)
        else:
            raise Exception("unsupported AE type: {}".format(kind))

    def call(self, inputs, training=None, mask=None):
        z = self.encoder(inputs)
        x_hat = self.decoder(z)
        return x_hat

    def train(self, X, batch_size, epochs, learning_rate, chart_file=""):
        optim = Adam(learning_rate=learning_rate)
        self.compile(optim, loss=losses.square_loss, metrics=metrics.relative_error)
        early_stop = EarlyStopping(monitor='val_loss', min_delta=1e-6, patience=20)
        history = self.fit(X, X, batch_size=batch_size, epochs=epochs, callbacks=[early_stop], validation_split=0.2)
        draw_train_history(history, chart_file=chart_file)
        # self.summary()


def main(argv):
    corpus = FLAGS.corpus
    data_dim = FLAGS.data_dim
    data_trans = FLAGS.data_trans
    kind = FLAGS.kind
    learning_rate = 1e-4
    batch_size = 256
    epochs = 100
    model_name = "ae"

    normal_data_files, anomaly_data_file, train_history_file, check_file, model_file, scaler_file = gen_outfile(
        data_dim,
        corpus,
        data_trans,
        model_name)
    X, ano_X = preprocess(normal_data_files, anomaly_data_file, data_dim)

    model = AE(kind=kind, input_shape=(X.shape[-2], X.shape[-1]), hidden_dims=[128, 32])
    model.train(X, batch_size=batch_size, epochs=epochs, learning_rate=learning_rate, chart_file=train_history_file)
    save_model(model, model_file)
    model = load_model(model_file)

    print("正常样本")
    e1 = fit_err_percentage(X, model.predict(X, batch_size=10000), losses.square_loss)
    print("故障样本")
    e2 = fit_err_percentage(ano_X, model.predict(ano_X, batch_size=10000), losses.square_loss)
    draw_fit_error(e1, e2, model_file)

    # 抽取一些正常和异常样本,看看编码器的拟合情况
    if data_dim == "1d":
        draw_fit_error_1d(X, ano_X, model(X), model(ano_X), chart_file=check_file)
    else:
        draw_fit_error_2d(X, ano_X, model(X), model(ano_X), chart_file=check_file)


if __name__ == '__main__':
    app.run(main)

# python .\anomaly_detection\ae.py --kind fc --corpus cwru --data_dim 1d --data_trans original
# python .\anomaly_detection\ae.py --kind fc --corpus jiangnan --data_dim 1d --data_trans original
